#include <iostream>
#include <array>
#include <ctime>

#ifdef WAFER_OPENSSL
#include <openssl/evp.h>
#endif

#include "container.h"

#define MSG_LEN 32
#define MSG_CNT 16
#define MD5_DIGEST_LEN 16L
#define INPUT_SIZE (MSG_LEN * MSG_CNT)
#define OUTPUT_SIZE (MD5_DIGEST_LEN * MSG_CNT)

#ifndef WAFER_OPENSSL
extern "C"  {
    void CUDA(MemRef<unsigned char, 1>* input, int n, MemRef<unsigned char, 1>* output, int N);
}

void MD5_OP(MemRef<unsigned char, 1>* input, int n, MemRef<unsigned char, 1>* output, int N) {
    CUDA(input, n, output, N);
}
#endif

uint8_t exampleInput[INPUT_SIZE];
MemRef<unsigned char, 1> input{exampleInput, INPUT_SIZE};

uint8_t *exampleOutput = new uint8_t[OUTPUT_SIZE];
MemRef<unsigned char, 1> output{exampleOutput, OUTPUT_SIZE};

void printResult(uint8_t* hash) {
    // Print the output.
    std::cout << "[ ";
    for (size_t i = 0; i < MD5_DIGEST_LEN; i++) {
      printf("%02x ", hash[i]);
    }
    std::cout << "]" << std::endl;
}

void printInput(int idx, uint8_t* input) {
  // Print the input.
  std::cout << "Input " << idx << ": [";
  for (size_t i = 0; i < MSG_LEN; i++) {
    std::cout << input[i];
  }
  std::cout << "] ==> ";
}

void generateResultGPUMD5() {
#ifdef WAFER_OPENSSL
  std::cout << "--------------------------------------------------------\nOpenSSL MD5\n";

  const clock_t begin_time = clock();
  const EVP_MD* md5 = EVP_md5();
  for (int i = 0; i < MSG_CNT; i++) {
    printInput(i, exampleInput + (i * MSG_LEN));

    EVP_MD_CTX *mdctx = EVP_MD_CTX_new();
    EVP_DigestInit_ex(mdctx, md5, NULL);
    EVP_DigestUpdate(mdctx, exampleInput + (i*MSG_LEN), MSG_LEN);
    uint32_t md_len = MD5_DIGEST_LEN;
    EVP_DigestFinal_ex(mdctx, exampleOutput + (i*MD5_DIGEST_LEN), &md_len);

    printResult(exampleOutput + (i * MD5_DIGEST_LEN));
  }
  float duration = float(clock() - begin_time) / CLOCKS_PER_SEC;
  std::cout << "Calculate duration: " << std::fixed << duration << std::endl;
  std::cout << "--------------------------------------------------------"
            << std::endl;

#else
  const clock_t begin_time = clock();

  #pragma wafer parallel
  MD5_OP(&input, MSG_LEN, &output, MSG_CNT);

  float duration = float(clock() - begin_time) / CLOCKS_PER_SEC;

  std::cout << "--------------------------------------------------------\nMLIR MD5\n";
  for (int i = 0; i < MSG_CNT; i++) {
    printInput(i, exampleInput + (i * MSG_LEN));
    printResult(exampleOutput + (i * MD5_DIGEST_LEN));
  }
  std::cout << "Calculate duration: " << std::fixed << duration << std::endl;
  std::cout << "--------------------------------------------------------"
            << std::endl;
#endif
}

void prepareInput() {
  for(int i = 0; i < INPUT_SIZE; i++) {
    exampleInput[i] = uint8_t(i % 26) + 'a';
  }
}

int main() {
  prepareInput();
  generateResultGPUMD5();
  return 0;
}
